package com.rizzo.shell.intern;


import com.google.common.collect.ImmutableSet;
import com.jcraft.jsch.*;
import com.rizzo.shell.intern.host.SshHost;
import com.rizzo.shell.intern.host.SshHost.Environment;
import com.rizzo.shell.intern.host.SshHost.InfoType;
import com.rizzo.shell.intern.host.SshHost.Type;
import org.apache.sshd.ClientChannel;
import org.apache.sshd.ClientSession;
import org.apache.sshd.SshClient;
import org.apache.sshd.client.future.AuthFuture;
import org.apache.sshd.client.future.ConnectFuture;
import org.apache.sshd.client.future.OpenFuture;
import org.springframework.shell.core.CommandMarker;
import org.springframework.shell.core.annotation.CliCommand;
import org.springframework.shell.core.annotation.CliOption;
import org.springframework.stereotype.Component;

import java.io.*;
import java.util.Set;
import java.util.logging.Level;
import java.util.logging.Logger;

@Component
public class SshCommands implements CommandMarker {

    public static final String LINUX_DISK_COMMAND = "df -h";
    public static final String LINUX_MEMORY_COMMAND = "free -mt";
    public static final String LINUX_PROCESS_COMMAND = "ps aux|grep java";
    public static final String LINUX_PORT_COMMAND = "netstat -tnl";
    public static final String LINUX_JAVA_COMMAND = "java -version";
    private static Logger logger = Logger.getLogger("com.rizzo.shell.intern.SshCommands");

	private Set<SshHost> sshLogins = ImmutableSet.of(
			new SshHost(Environment.ON, "multivac", 21, 10000L, "jalie", "Y5XOEUn6", Type.ACCOUNT,
                    "/var/log/syslog"),
            new SshHost(Environment.ON, "localhost", 21, 10000L, "jalie", "Y5XOEUn6", Type.ACCOUNT,
                    "/var/log/syslog"));

	@CliCommand(value = "ssh:execute", help = "Execute a command on one server with the ssh command")
	public String execute(
			@CliOption(key = {"host"}, mandatory = true, help = "The host to execute the command on") final String host,
			@CliOption(key = {"port"}, mandatory = true, help = "The host to execute the command on") final String port,
			@CliOption(key = {"timeout"}, mandatory = true, help = "The host to execute the command on") final String timeout,
			@CliOption(key = {"username"}, mandatory = true, help = "The username of the user") final String username,
			@CliOption(key = {"password"}, mandatory = true, help = "The password of the user") final String password,
			@CliOption(key = {"command"}, mandatory = true, help = "The command to execute") final String command) throws Exception {
		logger.info("Executing ssh command " + command + " ...");
		try {
			return executeCommand(host, Integer.parseInt(port), username, password, Long.parseLong(timeout), command);
		} catch (Exception e) {
			logger.log(Level.SEVERE, "Fail!", e);
		}
		return null;
	}

	@CliCommand(value = "servers:info", help = "Execute a command on a server with the ssh command")
	public String serverInfo(@CliOption(key = {"type"}, mandatory = true, help = "The type of server to query",
			unspecifiedDefaultValue = "ACCOUNT") final Type type,
			@CliOption(key = {"infoType"}, mandatory = false, help = "The command type to execute",
					unspecifiedDefaultValue = "ALL") final InfoType infoType,
			@CliOption(key = {"environment"}, mandatory = false, help = "The environment of the server",
					unspecifiedDefaultValue = "ON") final Environment environment,
            @CliOption(key = {"grepArg"}, mandatory = false, help = "Non mandatory grep argument for the log command",
                    unspecifiedDefaultValue = "") final String grepArg) throws Exception {

		for (SshHost sshLogin : sshLogins) {
			if (sshLogin.getType() == type && sshLogin.getEnvironment() == environment) {
				try {
					logger.log(Level.INFO, "                                                        ");
					logger.log(Level.INFO, "                                                        ");
					logger.log(Level.SEVERE, "################################ " + sshLogin.getUsername() + "@" + sshLogin.getHost() +
							" ########################################");
					logger.log(Level.INFO, "                                                        ");
					if (infoType == InfoType.DISK || infoType == InfoType.STATS || infoType == InfoType.ALL) {
						logger.log(Level.WARNING, "--------------------- DISK ------------------------");
						logger.log(Level.INFO, executeCommand(sshLogin.getHost(), sshLogin.getPort(), sshLogin.getUsername(),
                                sshLogin.getPassword(), sshLogin.getTimeOut(), LINUX_DISK_COMMAND));
					}
					if (infoType == InfoType.MEMORY || infoType == InfoType.STATS || infoType == InfoType.ALL) {
						logger.log(Level.WARNING, "--------------------- MEMORY ----------------------");
						logger.log(Level.INFO,
								executeCommand(sshLogin.getHost(), sshLogin.getPort(), sshLogin.getUsername(), sshLogin.getPassword(),
                                        sshLogin.getTimeOut(), LINUX_MEMORY_COMMAND));
					}
					if (infoType == InfoType.PROCESS || infoType == InfoType.STATS || infoType == InfoType.ALL) {
						logger.log(Level.WARNING, "--------------------- PROCESS ---------------------");
						logger.log(Level.INFO, executeCommand(sshLogin.getHost(), sshLogin.getPort(), sshLogin.getUsername(),
                                sshLogin.getPassword(), sshLogin.getTimeOut(), LINUX_PROCESS_COMMAND));
					}
					if (infoType == InfoType.PORTS || infoType == InfoType.STATS || infoType == InfoType.ALL) {
						logger.log(Level.WARNING, "--------------------- PORTS ---------------------");
						logger.log(Level.INFO, executeCommand(sshLogin.getHost(), sshLogin.getPort(), sshLogin.getUsername(),
                                sshLogin.getPassword(), sshLogin.getTimeOut(), LINUX_PORT_COMMAND));
					}
					if (infoType == InfoType.JAVA || infoType == InfoType.ALL) {
						logger.log(Level.WARNING, "--------------------- JAVA ----------------------");
						logger.log(Level.INFO, executeCommand(sshLogin.getHost(), sshLogin.getPort(), sshLogin.getUsername(),
                                sshLogin.getPassword(), sshLogin.getTimeOut(), LINUX_JAVA_COMMAND));
					}
					if (infoType == InfoType.LOG || infoType == InfoType.ALL) {
						logger.log(Level.WARNING, "--------------------- LOG ---------------------");
						logger.log(Level.INFO, executeCommand(sshLogin.getHost(), sshLogin.getPort(), sshLogin.getUsername(),
                                sshLogin.getPassword(), sshLogin.getTimeOut(),
								"cat " + sshLogin.getLogfileLocation() + " | egrep -i '" + grepArg + "'"));
					}
				} catch (Exception e) {
					logger.log(Level.SEVERE, "Failure!", e);
				}
			}
		}

		return "Done";
	}

	@CliCommand(value = "servers:exec", help = "Execute a command on a server with the ssh command")
	public String serverExec(@CliOption(key = {"type"}, mandatory = true, help = "The type of server to execute a command on") final Type type,
			@CliOption(key = {"environment"}, mandatory = false, help = "The environment of the server",
					unspecifiedDefaultValue = "ON") final Environment environment,
			@CliOption(key = {"command"}, mandatory = true, help = "The command to execute") final String command
	) throws Exception {
		for (SshHost sshLogin : sshLogins) {
			if (sshLogin.getType() == type && sshLogin.getEnvironment() == environment) {
				try {
					logger.log(Level.INFO, "                                                        ");
					logger.log(Level.INFO, "                                                        ");
					logger.log(Level.SEVERE, "################################ " + sshLogin.getUsername() + "@" + sshLogin.getHost() +
							" ########################################");
					logger.log(Level.INFO, "                                                        ");
					logger.log(Level.WARNING, "--------------------- COMMAND ---------------------");
					logger.log(Level.INFO, executeCommand(sshLogin.getHost(), sshLogin.getPort(), sshLogin.getUsername(),
                            sshLogin.getPassword(), sshLogin.getTimeOut(), command));
				} catch (Exception e) {
					logger.log(Level.SEVERE, "Fail!", e);
				}
			}
		}

		return "Done";
	}

	private String executeCommand(String host, int port, String username, String password, long timeOut, String command) throws Exception {
		byte[] result = null;
		SshClient client = SshClient.setUpDefaultClient();
		client.start();
		ConnectFuture connectFuture = client.connect(host, 22);
		connectFuture.await(timeOut);
		if (!connectFuture.isDone() || !connectFuture.isConnected()) {
			logger.log(Level.SEVERE, "Failed to connect to " + host + ":" + port + " within timeout " + timeOut + "ms");
		}
		AuthFuture authResult;
		ClientSession session = connectFuture.getSession();
		authResult = session.authPassword(username, password);
		authResult.await(timeOut);
		if (!authResult.isDone() || authResult.isFailure()) {
			logger.log(Level.SEVERE, "Failed to authenticate user " + username);
		}
		ClientChannel channel = session.createChannel(ClientChannel.CHANNEL_EXEC, command);
		ByteArrayInputStream in = new ByteArrayInputStream(new byte[]{0});
		channel.setIn(in);
		ByteArrayOutputStream out = new ByteArrayOutputStream();
		channel.setOut(out);
		ByteArrayOutputStream err = new ByteArrayOutputStream();
		channel.setErr(err);
		OpenFuture openFuture = channel.open();
		openFuture.await(timeOut);
		if (openFuture.isOpened()) {
			channel.waitFor(ClientChannel.CLOSED, 0);
			result = out.toByteArray();
		}
		client.stop();
		return new String(result);
	}

	protected Session getJschSession(final String host, final int port, final String username, final String password) throws JSchException {
		JSch sch = new JSch();
		Session session = sch.getSession(username, host, port);
		session.setUserInfo(new UserInfo() {
			public String getPassphrase() {
				return null;
			}

			public String getPassword() {
				return password;
			}

			public boolean promptPassword(String message) {
				return true;
			}

			public boolean promptPassphrase(String message) {
				return false;
			}

			public boolean promptYesNo(String message) {
				return true;
			}

			public void showMessage(String message) {
			}
		});
		session.connect();
		return session;
	}

	protected String readFile(String path, Session session) throws Exception {
		ChannelExec c = (ChannelExec) session.openChannel("exec");
		OutputStream os = c.getOutputStream();
		InputStream is = c.getInputStream();
		c.setCommand("scp -f " + path);
		c.connect();
		String header = readLine(is);
		int length = Integer.parseInt(header.substring(6, header.indexOf(' ', 6)));
		os.write(0);
		os.flush();

		byte[] buffer = new byte[length];
		length = is.read(buffer, 0, buffer.length);
		os.write(0);
		os.flush();

		c.disconnect();
		return new String(buffer);
	}

	private String readLine(InputStream in) throws IOException {
		ByteArrayOutputStream baos = new ByteArrayOutputStream();
		for (; ;) {
			int c = in.read();
			if (c == '\n') {
				return baos.toString();
			} else if (c == -1) {
				throw new IOException("End of stream");
			} else {
				baos.write(c);
			}
		}
	}

}
